#include <linux/if.h>
#include <linux/if_tun.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <stdarg.h>
#include <assert.h>

#define MAX_CPU_TYPE 16
#define MAX_IMAGE_PATH 256

#define TAP_DEV "tap0"
#define QUEUE_SIZE_MAX 16

#define MAX_ARG_OPTS 64
#define OPT_LEN 256

static int smp = 2;
// default to a53 for RK3399
static char *cpu = "cortex-a53";

static int mem = 1024;

static char *disk_image = NULL;

static char *kernel = NULL;

static int queue_size = 2;

static int uefi = 0;


char arg_opts[MAX_ARG_OPTS][OPT_LEN] = {0};


static char *gen_opt(char *fmt, ...)
{
    static int index = 0;
    char *buffer;

    assert(index < MAX_ARG_OPTS - 1);

    buffer = arg_opts[index++];

    va_list ap;
    va_start(ap, fmt);
    vsnprintf(buffer, OPT_LEN - 1,  fmt, ap);
    va_end(ap);

    return buffer;
}

// borrowed from Linux kernel documents tuntap.txt
int tun_alloc_mq(char *dev, int queues, int *fds)
{
    struct ifreq ifr;
    int fd, err, i;

    if (!dev)
        return -1;

    memset(&ifr, 0, sizeof(ifr));
    /* Flags: IFF_TUN   - TUN device (no Ethernet headers)
     *        IFF_TAP   - TAP device
     *
     *        IFF_NO_PI - Do not provide packet information
     *        IFF_MULTI_QUEUE - Create a queue of multiqueue device
     */
    ifr.ifr_flags = IFF_TAP | IFF_NO_PI | IFF_MULTI_QUEUE;
    strcpy(ifr.ifr_name, dev);

    for (i = 0; i < queues; i++) {
        if ((fd = open("/dev/net/tun", O_RDWR)) < 0)
           goto err;
        err = ioctl(fd, TUNSETIFF, (void *)&ifr);
        if (err) {
           perror("ioctl: ");
           close(fd);
           goto err;
        }
        fds[i] = fd;
    }

    return 0;
err:
    for (--i; i >= 0; i--)
        close(fds[i]);
    return err;
}

int prepare_and_fork_qemu()
{
    int tap_fds[QUEUE_SIZE_MAX] = {0};
    int i;
    int ret;
    char fds_arg[128] = {0};
    char number[8];

    char *args[64] = {NULL};

    assert(queue_size <= QUEUE_SIZE_MAX);

    printf("allocating %d queue on tap device %s\n", queue_size, TAP_DEV);
    ret = tun_alloc_mq(TAP_DEV, queue_size, tap_fds);

    if (ret < 0) {
        perror("tun_malloc_mq:");
        return ret;
    }

    printf("tap fd allocate succefully!\n");

    // generate fds arguments
    snprintf(fds_arg, 128, "fds=");
    for (i = 0; i < queue_size; i++) {
        snprintf(number, 8, "%d", tap_fds[i]);
        strcat(fds_arg, number);

        // add the ':', except last one
        if (i < queue_size - 1) {
            strcat(fds_arg, ":");
        }
    }

    // set up arguments
    i = 0;
    args[i++] = "taskset";
    args[i++] = "f"; // cpu: 0, 1, 2, 3
    args[i++] = "./qemu-system-aarch64";
    args[i++] = "-M";
    args[i++] = "virt";
    args[i++] = "-smp";
    args[i++] = gen_opt("%d", smp);
    args[i++] = "-cpu";
    args[i++] = cpu;
    args[i++] = "-m";
    args[i++] = gen_opt("%d", mem);
    args[i++] = "-enable-kvm";

    if (!uefi) { // with -kernel and -append
        args[i++] = "-kernel";
        args[i++] = kernel;
        args[i++] = "-append";
        args[i++] = "'rootfstype=ext4 root=/dev/vda2 console=ttyAMA0 rw net.ifnames=0'";
    }
    else {
        args[i++] = "-pflash";
        args[i++] = "flash0.img";
    }

    args[i++] = "-nographic";
    args[i++] = "-drive";
    args[i++] = gen_opt("file=%s,if=none,id=disk0,format=raw", disk_image);
    args[i++] = "-device";
    args[i++] = "virtio-blk-pci,drive=disk0,scsi=off,bootindex=1";
    args[i++] = "-device";
    args[i++] = "virtio-net-pci,mac=9a:56:57:58:59:5a,netdev=net0,mq=on,vectors=10,rx_queue_size=1024,tx_queue_size=1024";
    args[i++] = "-netdev";
    args[i++] = gen_opt("tap,vhost=on,id=net0,%s", fds_arg);

    // printf cmdline
    i = 0;
    printf("\nFinal qemu cmdline:\n");
    while (args[i] != NULL) {
        printf("%s\n", args[i++]);
    }

    printf("\n\n");

    // execl qemu
    execvp(args[0], args);

    return 0;
}

void usage(char **argv)
{
    fprintf(stderr, "********************************************************************************\n");
    fprintf(stderr, "   a wrapper for qemu to support multiple queue tuntap and virtio/vhost net\n");
    fprintf(stderr, "********************************************************************************\n");
    fprintf(stderr, "Usage: %s options  \n", argv[0]);
    fprintf(stderr, "options:  \n");
    fprintf(stderr, "-d: qemu disk image path \n");
    fprintf(stderr, "-m: virtual machine memory size in Mb (default 1020M)\n");
    fprintf(stderr, "-c: qemu -cpu argument (default cortex-a53)\n");
    fprintf(stderr, "-s: vcpu number (default 2 cpu)\n");
    fprintf(stderr, "-k: the linux kernel image path\n");
    fprintf(stderr, "-q: the queue size (default 2)\n");
    fprintf(stderr, "-u: start qemu with uefi (default flash0.img)\n");
    fprintf(stderr, "-h: show this help message\n");
}

int main(int argc, char **argv)
{
    int ch;
    opterr = 0;
    char *pend;
    while ((ch = getopt(argc, argv, "d:m:c:s:k:q:uh")) != -1) {
        switch (ch) {
            case 'd':
                disk_image = optarg;
                break;
            case 'm':
                mem = strtol(optarg, &pend, 10);
                break;
            case 'c':
                cpu = optarg;
                break;
            case 's':
                smp = strtol(optarg, &pend, 10);
                break;
            case 'k':
                kernel = optarg;
                break;
            case 'q':
                queue_size = strtol(optarg, &pend, 10);
                break;
            case 'u':
                uefi = 1;
                break;

            case 'h':
                usage(argv);
                exit(0);

            default:
                fprintf(stderr, "Error: unknown option: -%c\n", optopt);
                exit(-1);
        }

    }

    if (disk_image == NULL) {
        fprintf(stderr, "qemu disk image path must provide!\n");
        fprintf(stderr, "please run %s -h for help\n", argv[0]);
        exit(-1);
    }

    if (!uefi) {
        if (kernel == NULL) {
            fprintf(stderr, "on non-UEFI, linux kernel image path must provide!\n");
            fprintf(stderr, "please run %s -h for help\n", argv[0]);
            exit(-1);
        }
    }

    prepare_and_fork_qemu();

    return 0;
}
